'''
pyrwi.tools - General functions that don't fit anywhere else
'''
import scipy as sp

########################################################################
# Binary Operations
########################################################################
def getBits(word, bits, reverse=False) :
    '''
    Provides array support for getBit
    @param word: Integer to query
        NOTE: If word is negative, sign magnitude format is used.
    @param bit: Number of bit to check. List or integer. 0 based from LSB
    @return: If bits is a list, then returns a list. Else returns a 1 or 0
    
    Example:
    getBits(7, range(4)) #Converts 7 to binary with LSB at element 0
    getBits(7, range(4), True) #Converts 7 to binary with MSB at element 0
    '''
    if not hasattr(bits, '__iter__') :
        #If bits is a single integer, return a single integer
        return (word>>bits)&1

    msb = len(bits)-1
    rtn = [0]*(msb+1)

    for i,b in enumerate(bits) :
        rtn[i] = (word>>b) & 1 #shift to bit, I want & take only lowest bit

    if reverse: rtn.reverse()
    return rtn

def binWeight(agg, bit) :
    '''
    reduce() callback for bitsToInt. msbFirst=True always
    @param agg: Binary aggregate so far
    @param bit: next bit to add to aggregate
    @return Integer that is binary weighted sum of array
    '''
    return (agg<<1)+bit
def bitsToInt(bits, msbFirst=False) :
    '''
    Take an array of 1's and 0's and compute an integer from them
    @param bits: Array of 1's and 0's
    @param msbFirst: If true bits[0] is most significant bit
    @return integer
    '''
    if not msbFirst: 
        bits = bits[:] #Copy elements so that I don't change caller's data
        bits.reverse()
        
    return reduce(binWeight, bits)

########################################################################
# Formatting helpers
########################################################################
def fmtEng( val, digits=None):
    '''
    Format 'val' in engineering notation.
    @param val : Value to format
    @param digits : Number of digits after decimal
    @return formatted string
    '''
    import math

    if digits is None :
        fmt = '%fe%d'
    else :
        fmt = '%%.%dfe%%d' % digits

    #Get exponent
    l = math.log10(abs(val))
    l = math.floor(l)
    p3 = int(int(l) / 3) * 3 #Round to nearest 3

    multiplier = val / pow(10, p3)
    return fmt % (multiplier, p3)

def getEngExp(val) :
    '''
    Find the 10's exponent and round to the nearest 3
    @param val : Floating point or Int input
    @return: One of ... -6, -3, 0, 3, 6, 9 ...
    '''
    import math
    l = math.log10(abs(val))
    l = math.floor(l)
    return int(int(l) / 3) * 3 #Round to nearest 3
    
UnitsSpice = {-15:'f', -12:'p', -9:'n', -6:'u', -3:'m', 0:'',
             3:'k', 6:'M', 9:'G'}

def fmtBin(x, digits=0):
    '''
    Take a number and generate a binary format in a string
    @param x (int): Number to convert or array of numbers
    @param digits : Number of binary digits
    Invert with: int(ut.bin(255),2)
    '''
    #Handle arrays
    if hasattr(x, '__iter__') :
        return [ bin(v) for v in x]
    oct2bin = ['000','001','010','011','100','101','110','111']
    binstring = [oct2bin[int(n)] if n != '-' else '-' for n in oct(x)]
    return ''.join(binstring).lstrip('0').zfill(digits)

########################################################################
#  helpers
########################################################################
def writeDataFile(outFile, *data, **kargs) :
    '''
    Write lists of data to a file formatted nicely for Excel.
    @param outFile: Output file name or object
    @param headers: List of column headers
        NOTE: len(headers) must == len(data)
    @param data: List of lists of data
        NOTE: All data will be formatted as floating point
    
    Usage:
    writeDataFile('log.csv', [1,2,3,4], [4,5,1,3])
    writeDataFile('log.csv', [[1,2,3,4], [4,5,1,3]])
    writeDataFile(open('log.csv','a'), [1,2,3,4], [4,5,1,3])
    writeDataFile('log.csv', headers=['time','Volts'], range(4), [4,5,1,3])
    '''
    from copy import deepcopy
    
    if isinstance(outFile, file) : 
        #Already a file handle
        fh = outFile
        closeFile = False
    else :
        #Open a file, default to overwriting any existing file
        fh = open(outFile, 'w')
        closeFile = True
    
    if kargs.has_key('headers') :
        headers = kargs['headers']
        if len(headers) != len(data): raise IndexError, "Headers must be same length as data"
        fh.write( ",".join([str(h) for h in headers]) + "\n")
        
    #Make all vectors the same length
    maxLen = len(max(data, key=len)) #Max returns longest vector
    data = [ v.tolist() if hasattr(v, "tolist") else deepcopy(v) for v in data]
    for v in data :
        v += [''] * (maxLen - len(v))
    
    #Make a pattern for writting one row at a time
    pat = ",".join(["%f"]*len(data)) + "\n"
    
    if True :
        import csv
        csvWrite = csv.writer(fh, lineterminator="\n")
        for tpl in zip(*data) :
            csvWrite.writerow(tpl)
    else :
        for i,tpl in enumerate(zip(*data)) :
            fh.write( pat % tpl )
    
    if closeFile: fh.close()
    
########################################################################
# Signals & Signal Processing
########################################################################
def timeVector(Fs, nSamples) :
    return sp.arange(0,nSamples/Fs, step=1.0/Fs)

def analog_fft(Fs, a) :
    '''
    Compute the FFT of an analog (ie. consider time info) signal
    @param Fs: Sampling frequency
    @param a: Vector of time domain data
    @return (frequency points, fft(a))
    '''
    import numpy.fft
    
    fa = numpy.fft.rfft(a) #Use real FFT, because I assume my data is real
    fa = fa / len(fa)
    freq = sp.linspace(0, Fs/2.0, num=len(fa))
    
    return (freq, fa)
    
